from AStock.ASEvent import Worker
from AStock.ASSetting import settings, cons
from AStock.ASUtil import date_to_bar_format, date_from_bar_format, ASHttpParser
import json
from pymongo import UpdateOne


class ASResearchReportCrawler(Worker):

    def __init__(self):
        super(ASResearchReportCrawler, self).__init__()
        self._proxy = None
        self._handler = None
        self._date_mapping = None
        self._fields = []
        if settings.research_report_source == cons.CrawlSource.EASTMONEY:
            self._url = cons.EASTMONEY_RESEARCH_REPORT
            self._date_mapping = self._eastmoney_date
            self._handler = self._eastmoney_research_report_handler
            self._fields = [
                "title",
                "stockName",
                "stockCode",
                "orgSName",  # 机构名称
                "publishDate",  # 发布日期
                # 标识码，如AP202009171413709566，
                # 可定位到研报pdf的url
                # http://pdf.dfcfw.com/pdf/H3_AP202009171413709566_1.pdf
                "infoCode",
                "predictNextTwoYearEps",
                "predictNextTwoYearPe",
                "predictNextYearEps",
                "predictNextYearPe",
                "predictThisYearEps",
                "predictThisYearPe",
                "indvInduName",  # 所属行业
                "count"  # 近1月个股研报数
            ]

    @staticmethod
    def _eastmoney_date(date):
        return date_to_bar_format(date)

    @staticmethod
    def _eastmoney_research_report_handler(text, **kwargs):
        """
        解析响应的研报内容
        :param text: str，待解析的响应文本
        :return: dict，
            失败时，None
            成功时，{'page_no': 1, 'pages': 105, 'data': [report1, report2, ...]}
            其中，reportx为股票的研报数据，格式为：
            {
                "title": "销售持续恢复，拿地趋于积极",
                "stockName": "保利地产",
                "stockCode": "600048",
                "orgSName": "华西证券",
                "publishDate": "2020-09-13 00:00:00",
                "infoCode": "AP202009131411540332",
                "predictNextTwoYearEps": "3.8800000000",
                "predictNextTwoYearPe": "4.3000000000",
                "predictNextYearEps": "3.3200000000",
                "predictNextYearPe": "5.0000000000",
                "predictThisYearEps": "2.8300000000",
                "predictThisYearPe": "5.9000000000",
                "indvInduName": "房地产",
                "count": 24
            }
        """
        start_pos = text.find('{')
        end_pos = text.rfind('}')
        if start_pos == -1 or end_pos == -1:
            print('eastmoney research report text format error')
            return None
        charset = kwargs.get('charset', 'utf-8')
        try:
            js = json.loads(text[start_pos:end_pos + 1], encoding=charset)
        except json.JSONDecodeError as e:
            print('error parse eastmoney research report: {}'.format(e.msg))
            return None

        try:
            fields = kwargs.get('fields', [])
            res = {'page_no': kwargs.get('pageno'), 'pages': None, 'data': []}
            if js['data']:
                data = js['data']
                res['pages'] = js['TotalPage']
                for item in js['data']:
                    report = {k: v for k, v in item.items() if k in fields}
                    res['data'].append(report)
                return res
        except ValueError as e:
            print('error parse eastmoney financial notice date: {}'.format(e))
            return None
        return None

    def run(self, event_data):
        """
        执行抓取研报的任务
        :param event_data: dict，包括：
            event_type: research_report 抓取研报
            begin_date: str，研报开始日期，格式为 yyyymmdd
            end_date: str，研报结束日期，格式为 yyyymmdd
            page_no: int，页码
            page_size: int，页大小
            rand: int，随机数
            timestamp: int，时间戳

        :return: dict，
            失败时，返回的字典为 {'error': 'error string'}
            成功时，{'page_no': 1, 'pages': 105, 'data': [report1, report2, ...]}
            其中，reportx为股票的研报数据，格式为：
            {
                "title": "销售持续恢复，拿地趋于积极",
                "stockName": "保利地产",
                "stockCode": "600048",
                "orgSName": "华西证券",
                "publishDate": "2020-09-13 00:00:00",
                "infoCode": "AP202009131411540332",
                "predictNextTwoYearEps": "3.8800000000",
                "predictNextTwoYearPe": "4.3000000000",
                "predictNextYearEps": "3.3200000000",
                "predictNextYearPe": "5.0000000000",
                "predictThisYearEps": "2.8300000000",
                "predictThisYearPe": "5.9000000000",
                "indvInduName": "房地产",
                "count": 24
            }
        """
        try:
            event_type = event_data['event_type']
            begin_date = event_data['begin_date']
            end_date = event_data['end_date']
            page_no = event_data['page_no']
            page_size = event_data['page_size']
            rand = event_data['rand']
            timestamp = event_data['timestamp']
        except KeyError as e:
            return {'error': 'key error in event data, key: {}'.format(e)}

        if event_type not in ['research_report']:
            return {'error': 'event type error: {}'.format(event_type)}

        if self._date_mapping:
            begin_date = self._date_mapping(begin_date)
            end_date = self._date_mapping(end_date)

        url_params = {
            'pageno': page_no,
            'begin_date': begin_date,
            'end_date': end_date,
            'pagesize': page_size,
            'rand': rand,
            'timestamp': timestamp
        }
        url = self._url.format(**url_params)
        parser = ASHttpParser(url, cons.HEADER, self._proxy, content_type="text")
        res = parser.parse_text(
            text_handler=self._handler,
            charset=parser.charset,
            fields=self._fields,
            **url_params
        )
        if res is None:
            return {'error': 'error crawl research report, url: {}'.format(url)}
        return res


class ASResearchReportSaver(Worker):

    def __init__(self):
        super().__init__()

    def run(self, event_data):
        """
        执行保存研报的任务
        :param event_data: dict，包括：
            event_type: research_report 保存研报
            data: list，[report1, report2, ...]
                其中，reportx为股票的研报数据，格式为：
                {
                    "title": "销售持续恢复，拿地趋于积极",
                    "stockName": "保利地产",
                    "stockCode": "600048",
                    "orgSName": "华西证券",
                    "publishDate": "2020-09-13 00:00:00",
                    "infoCode": "AP202009131411540332",
                    "predictNextTwoYearEps": "3.8800000000",
                    "predictNextTwoYearPe": "4.3000000000",
                    "predictNextYearEps": "3.3200000000",
                    "predictNextYearPe": "5.0000000000",
                    "predictThisYearEps": "2.8300000000",
                    "predictThisYearPe": "5.9000000000",
                    "indvInduName": "房地产",
                    "count": 24
                }

        :return: dict，
            成功：{'succ': 'saved: research report'}
            失败：{'fail': 'failed to save research report'}
        """
        try:
            event_type = event_data['event_type']
            data = event_data['data']
        except KeyError as e:
            return {'fail': 'key error in event data, key: {}'.format(e)}

        if event_type not in ['research_report']:
            return {'fail': 'event type error: {}'.format(event_type)}

        succ = 'saved: research report'
        fail = 'failed to save research report'

        bulk = []
        for item in data:
            update_one = UpdateOne(
                # 用stockCode和infoCode作为匹配条件
                {'stockCode': item['stockCode'], 'infoCode': item['infoCode']},
                {'$set': self._reformat_data(item)},
                upsert=True
            )
            bulk.append(update_one)

        try:
            coll = settings.database.research_report
            coll.bulk_write(bulk)
        except Exception as e:
            return {'fail': fail + '    ' + str(e)}

        return {'succ': succ}

    def _reformat_data(self, report):
        res = {}
        for k, v in report.items():
            # print(k, " ", v)
            if k == 'publishDate':
                # "publishDate": "2020-09-13 00:00:00"
                # 调整为 20200913
                res[k] = int(date_from_bar_format(v[:10]))
            elif k in ['predictNextTwoYearEps', 'predictNextYearEps', 'predictThisYearEps']:
                try:
                    res[k] = round(float(v), 3)
                except ValueError:
                    res[k] = None
            elif k in ['predictNextTwoYearPe', 'predictNextYearPe', 'predictThisYearPe']:
                try:
                    res[k] = round(float(v), 2)
                except ValueError:
                    res[k] = None
            else:
                res[k] = v
        return res
